package com.llyb.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

@RestController
@Slf4j
public class FunctionChatModelController {

    @Resource
    private EmbeddingModel embeddingModel;

    @Resource
    private QdrantClient qdrantClient;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;


    // http://localhost:9012/embedding/embed
    @GetMapping("/embedding/embed")
    public String embed() {

        String prompt = """
                快跑，速度要快！
                """;
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);

        System.out.println(embeddingResponse);

        return embeddingResponse.content().toString();
    }

    @GetMapping("/embedding/collection")
    public void createCollection() {

        var vectorParam = Collections.VectorParams.
                newBuilder().
                setDistance(Collections.Distance.Cosine).
                setSize(1024).
                build();

        qdrantClient.createCollectionAsync("test-qdrant",vectorParam);
    }

    @GetMapping("/embedding/add")
    public String add() {

        String prompt = "以坚定信念前行。";
        TextSegment segment1 = TextSegment.from(prompt);
        segment1.metadata().put("author","llyb");
        Embedding content = embeddingModel.embed(segment1).content();
        String result = embeddingStore.add(content, segment1);

        System.out.println(result);
        return result;
    }

    @GetMapping("/embedding/query1")
    public String query1() {

        Embedding content = embeddingModel.embed("跑的很快").content();

        EmbeddingSearchRequest embeddingSearchRequest  = EmbeddingSearchRequest.builder().
                                                queryEmbedding(content).
                                                maxResults(3).
                                                build();

        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);

        System.out.println(searchResult.matches().get(0).embedded().text());
        System.out.println(searchResult.matches().get(1).embedded().text());
        System.out.println(searchResult.matches().get(2).embedded().text());
        return searchResult.matches().get(0).embedded().text();
    }

    @GetMapping(value = "/embedding/query2")
    public void query2(){
        Embedding queryEmbedding = embeddingModel.embed("慢").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(metadataKey("author").isEqualTo("llyb"))
                .maxResults(1)
                .build();

        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);

        System.out.println(searchResult.matches().get(0).embedded().text());
    }




}
